import argparse
import os
from glob import glob

from sklearn.model_selection import train_test_split

parser = argparse.ArgumentParser()
parser.add_argument('--data_root', help='dataset', default='lrs2_preprocessed', type=str)
args = parser.parse_args()


def trained_data_name_format(base_path):
    result = list(glob("{}/*".format(base_path)))
    result_list = []
    for i, dirpath in enumerate(result):
        dirs = os.listdir(dirpath)
        for f in dirs:
            replace = dirpath.replace(base_path, '').replace('/', '').replace('\\', '')
            result_list.append(str(os.path.join(replace, f)))
            if len(result_list) < 14:
                test_result = val_result = train_result = result_list
            else:
                train_result, test_result = train_test_split(result_list, test_size=0.2, random_state=42)
                test_result, val_result = train_test_split(test_result, test_size=0.5, random_state=42)
            for file_name, dataset in zip(("train.txt", "test.txt", "val.txt"),
                                          (train_result, test_result, val_result)):
                with open(os.path.join("filelists", file_name), 'w', encoding='utf-8') as fi:
                    fi.write("\n".join(dataset))


if __name__ == '__main__':
    path = args.data_root
    trained_data_name_format(path)
